use std::{
    fmt,
    fmt::Debug,
    ops::AddAssign,
    ops::Mul,
    ops::Add,
    sync::mpsc,
    thread};

use anyhow::anyhow;

use crate::Vector;

const NUM_THREADS: usize = 4;

pub struct Matrix<T> {
    pub data: Vec<T>,
    pub row: usize,
    pub col: usize,
}

struct MsgInput<T> {
    idx: usize,
    row: Vector<T>,
    col: Vector<T>,
}

struct MsgOutput<T> {
    idx: usize,
    value: T,
}

struct Msg<T> {
    input: MsgInput<T>,
    // sender to send the result back
    sender: oneshot::Sender<MsgOutput<T>>,
}

pub fn do_product<T>(a: Vector<T>, b: Vector<T>) -> anyhow::Result<T>
    where T: Add<Output=T> + Mul<Output=T> + AddAssign + Copy + Debug + Default {
    if a.len() != b.len() {
        return Err(anyhow!("Matrix dimensions do not match"));
    }

    let mut result: T = T::default();
    for i in 0..a.len() {
        result += a[i] * b[i];
    }
    Ok(result)
}

pub fn multiply<T>(a: &Matrix<T>, b: &Matrix<T>) -> anyhow::Result<Matrix<T>>
    where T: Add<Output=T> + Mul<Output=T> + AddAssign + Copy + Debug + Default + Send + 'static {
    if a.col != b.row {
        return Err(anyhow!("Matrix dimensions do not match"));
    }

    let senders = (0..NUM_THREADS).map(|_| {
        let (sender, receiver) = mpsc::channel::<Msg<T>>();

        thread::spawn(move || {
            for msg in receiver {
                let value = do_product(msg.input.row, msg.input.col)?;
                if let Err(e) = msg.sender.send(MsgOutput {
                    idx: msg.input.idx,
                    value,
                }) {
                    eprintln!("Send error: {:?}", e);
                }
            }

            Ok::<_, anyhow::Error>(())
        });

        sender
    }).collect::<Vec<_>>();

    let matrix_len = a.row * b.col;
    let mut data = vec![T::default(); matrix_len];
    let mut receivers = Vec::with_capacity(matrix_len);

    for i in 0..a.row {
        for j in 0..b.col {
            let row = Vector::new(&a.data[i * a.col..(i + 1) * a.col]);
            let col_data = b.data[j..].iter().step_by(b.col).copied().collect::<Vec<_>>();
            let col = Vector::new(col_data);
            let idx = i * b.col + j;

            let msg_input = MsgInput::new(idx, row, col);
            let (tx, rx) = oneshot::channel();
            let msg = Msg::new(msg_input, tx);
            if let Err(e) = senders[idx % NUM_THREADS].send(msg) {
                eprintln!("Send error: {:?}", e);
            }
            receivers.push(rx);
        }
    }

    for rx in receivers {
        let output = rx.recv()?;
        data[output.idx] = output.value;
    }

    Ok(Matrix {
        data,
        row: a.row,
        col: b.col,
    })
}

impl<T: Debug> Matrix<T> {
    pub fn new(data: impl Into<Vec<T>>, rows: usize, cols: usize) -> Matrix<T> {
        Self {
            data: data.into(),
            row: rows,
            col: cols,
        }
    }
}

impl<T> fmt::Display for Matrix<T>
    where T: fmt::Display {
    // display a 2x3 as {1 2 3, 4 5 6}, 3x2 as {1 2, 3 4, 5 6}
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "{{")?;
        for i in 0..self.row {
            for j in 0..self.col {
                write!(f, "{}", self.data[i * self.col + j])?;
                if j != self.col - 1 {
                    write!(f, " ")?;
                }
            }

            if i != self.row - 1 {
                write!(f, ", ")?;
            }
        }
        write!(f, "}}")?;
        Ok(())
    }
}

impl<T> Debug for Matrix<T>
    where
        T: fmt::Display,
{
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "Matrix(row={}, col={}, {})", self.row, self.col, self)
    }
}

impl<T> MsgInput<T> {
    fn new(idx: usize, row: Vector<T>, col: Vector<T>) -> MsgInput<T> {
        Self { idx, row, col }
    }
}

impl<T> Msg<T> {
    fn new(input: MsgInput<T>, sender: oneshot::Sender<MsgOutput<T>>) -> Msg<T> {
        Self { input, sender }
    }
}

impl<T> Mul for Matrix<T>
    where T: Add<Output=T> + Mul<Output=T> + AddAssign + Copy + Debug + Default + Send + 'static {
    type Output = Self;

    fn mul(self, rhs: Self) -> Self::Output {
        multiply(&self, &rhs).expect("Matrix multiply error")
    }
}